package com.iakuil.app.common.core.util;

import cn.hutool.core.lang.tree.Tree;
import cn.hutool.core.lang.tree.TreeNodeConfig;
import cn.hutool.core.lang.tree.TreeUtil;
import lombok.experimental.UtilityClass;
import org.springframework.cglib.beans.BeanMap;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树状结构工具类
 *
 * @author Kai
 */
@UtilityClass
public class TreeUtils {

    /**
     * 默认的树配置
     */
    private static final TreeNodeConfig DEFAULT_CONFIG = createConfig("id", "pid", "sort", null, null);


    /**
     * 将源数据转化为树结构数据
     *
     * @param originalList 数据集
     * @param rootId       根节点id
     * @param config       配置项
     * @return 数据集
     */
    public static <T, E> List<Tree<E>> build(List<T> originalList, E rootId, TreeNodeConfig config) {
        if (CollectionUtils.isEmpty(originalList)) {
            return Collections.emptyList();
        }

        // 转树，注意：Tree<E>里面泛型为ID的类型
        return TreeUtil.build(originalList, rootId, config, (object, tree) -> {
            Map<String, T> map = BeanMap.create(object);
            for (Map.Entry<String, T> entry : map.entrySet()) {
                tree.putExtra(entry.getKey(), entry.getValue());
            }
        });
    }

    /**
     * 将源数据转化为树结构数据
     *
     * @param originalList 数据集
     * @param rootId       根节点ID
     * @return 数据集
     */
    public static <T, E> List<Tree<E>> build(List<T> originalList, E rootId) {
        return build(originalList, rootId, DEFAULT_CONFIG);
    }

    /**
     * 查询所有祖先节点（包括自己）
     *
     * @param originalList   数据集
     * @param currentNode    当前节点对象
     * @param idMapper       当前节点对象ID函数
     * @param parentIdMapper 父节点对象ID函数
     * @return 对象列表
     */
    public static <T, K> List<T> getAncestor(List<T> originalList, T currentNode, Function<? super T, ? extends K> idMapper, Function<? super T, ? extends K> parentIdMapper) {
        if (CollectionUtils.isEmpty(originalList)) {
            return Collections.emptyList();
        }

        // 此集合用来接收实体对象
        final Set<T> nodes = new HashSet<>();
        // 将源对象转化为以ID为Key的Map
        Map<? extends K, T> entityMap = originalList.stream().collect(
            Collectors.toMap(idMapper, Function.identity(), (v1, v2) -> v2)
        );
        getParent(nodes, currentNode, entityMap, parentIdMapper);
        return new ArrayList<>(nodes);
    }

    /**
     * 节点查找子节点（包括自己）
     *
     * @param originalList 数据集
     * @param id           当前节点id
     * @param rootId       根节点id
     * @return 树结构
     */
    public static <T, E> Tree<E> getChildren(List<T> originalList, E id, E rootId) {
        return getChildren(build(originalList, rootId), id);
    }

    /**
     * 查询所有子节点（包括自己）
     *
     * @param trees  数据集
     * @param rootId 根节点id
     * @return 树结构
     */
    public static <E> Tree<E> getChildren(List<Tree<E>> trees, E rootId) {
        Tree<E> childNode = new Tree<>();
        if (CollectionUtils.isNotEmpty(trees)) {
            for (Tree<E> node : trees) {
                childNode = TreeUtil.getNode(node, rootId);
                if (Objects.nonNull(childNode)) {
                    break;
                }
            }
        }
        return childNode;
    }

    /**
     * 查找父节点
     *
     * @param nodeList       实体列表
     * @param currentNode    当前节点
     * @param entityMap      以ID为Key的Map
     * @param parentIdMapper 父节点ID函数
     */
    private static <T, K> void getParent(Set<T> nodeList, T currentNode, Map<? extends K, T> entityMap, Function<? super T, ? extends K> parentIdMapper) {
        if (Objects.nonNull(currentNode)) {
            nodeList.add(currentNode);
        }
        if (Objects.nonNull(currentNode)) {
            getParent(nodeList, entityMap.get(parentIdMapper.apply(currentNode)), entityMap, parentIdMapper);
        }
    }

    /**
     * 默认的树状构建配置
     *
     * @param idKey       ID属性名
     * @param parentIdKey 父节点ID属性名
     * @param weighKey    权重属性名
     * @param deep        解析的最大深度
     * @return 树的配置项
     */
    private static TreeNodeConfig createConfig(String idKey, String parentIdKey, String weighKey, String childrenKey, Integer deep) {
        // 配置
        TreeNodeConfig config = new TreeNodeConfig();
        if (StringUtils.isNotBlank(idKey)) {
            config.setIdKey(idKey);
        }
        if (StringUtils.isNotBlank(parentIdKey)) {
            config.setParentIdKey(parentIdKey);
        }
        if (StringUtils.isNotBlank(weighKey)) {
            config.setWeightKey(weighKey);
        }
        if (StringUtils.isNotBlank(childrenKey)) {
            config.setChildrenKey(childrenKey);
        }
        if (deep != null) {
            config.setDeep(deep);
        }
        return config;
    }
}


